#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <openssl/sha.h>
#include <pwd.h>


#include "PublicChain.h"
#include "cjson/cJSON.h"

#include "teex/base64.h"
#include "teex/util.h"
#include "teex/error.h"
#include "teex/task.h"
#include "teex/const.h"
#include "teex/debug.h"



int checkTaskInfo(teexTask *taskInfo)
{
	int i = 0;
	if (taskInfo == NULL)
		return -1;

	if (!isNumberStr(taskInfo->serviceID) || strlen(taskInfo->serviceID) != ID_LEN)
		return -1;

	if (taskInfo->input == NULL)
		return -1;
	
	if (taskInfo->dataNumber != 0 && taskInfo->dataList == NULL)
		return -1;

	for (i = 0; i < taskInfo->dataNumber; i++) {

		if (!isNumberStr(taskInfo->dataList[i]) || strlen(taskInfo->dataList[i]) != ID_LEN)
			return -1;
	}
	
	if (!isNumberStr(taskInfo->publicKey) || strlen(taskInfo->publicKey) != PUBLIC_KEY_LEN)
		return -1;
	
	if (!isNumberStr(taskInfo->price))
		return -1;

	return 0;
}

teex_status_t createTaskOnChain(char *chainAddr, int chainPort, char *privateKey, char *publicKey,
									char *serviceContractAddr, char *tokenContractAddr,
									teexTask *taskInfo, char **taskID,
									unsigned int timeout)
{

	int retval = TEEX_ERROR_UNIMPLEMENTED, len = 0, r = 0;
	char *hash_buf;
	char chainURL[1024];

	/* check all the arguments */
	if (chainAddr == NULL)
		return TEEX_ERROR_INVALID_CHAINURL;
	if (privateKey == NULL)
		return TEEX_ERROR_INVALID_PRIVATE_KEY;
	if (publicKey == NULL)
		return TEEX_ERROR_INVALID_PUBLIC_KEY;
	if (serviceContractAddr == NULL)
		return TEEX_ERROR_INVALID_CONTRACT1;
	if (tokenContractAddr == NULL)
		return TEEX_ERROR_INVALID_CONTRACT2;
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (taskID == NULL)
		return TEEX_ERROR_INVALID_TASKID_BUF;
	
	memset(chainURL, 0, 1024);
	sprintf(chainURL, "http://%s:%d", chainAddr, chainPort);

	/* connect to the blockchain */
	PublicChain pc = PublicChain(chainURL, serviceContractAddr, tokenContractAddr,
									privateKey, publicKey);
	Task task, task_res;

	if (pc.status != 0) {
		return TEEX_ERROR_CHAIN_DISCONNECTED;
	}

	/* check the taskInfo */
	if (checkTaskInfo(taskInfo) < 0) {
		return TEEX_ERROR_INVALID_TASKINFO;
	}

	/* create a onchain task */
	task.data_num = taskInfo->dataNumber;
	task.bound_data_list = taskInfo->dataList;
	task.servideID = taskInfo->serviceID;

	/* TODO: currently timeout is useless */
	task.timeout = 10;

	len = strlen(taskInfo->price);
	memset(task.payment, '0', 64-len);
	strncpy(task.payment+64-len, taskInfo->price, len);
	strncpy(task.pk_user, taskInfo->publicKey, PUBLIC_KEY_LEN);

	/* calculate the input hash */
	hash_buf = sha256AndToHexStr(taskInfo->input, strlen(taskInfo->input));
	if (hash_buf == NULL) {
		return TEEX_ERROR_INVALID_TASKINFO;
	}
	memcpy(task.input_hash, hash_buf, 2*HASH_LEN);
	

	/* Send the task to chain */
	if ((r = pc.newTask(&task)) != 0) 
		return TEEX_ERROR_CREATE_TASK;
	
	/* confirm the task onchain */
	while (timeout > 0 && (r=pc.getTaskByID(task.taskID, task_res)) != 0) {
		if (r != 5) 
			return TEEX_ERROR_CONFIRM;

		sleep(1);
		timeout --;
	}

	if (timeout <= 0 || r != 0) 
		return TEEX_ERROR_TIMEOUT;
		
	
	/* return the taskID */
	*taskID = (char*)malloc(ID_LEN+1);
	if (*taskID == NULL)
		return TEEX_ERROR_OUT_MEMORY;

	memset(*taskID, 0, ID_LEN+1);
	strncpy(*taskID, task.taskID, ID_LEN);

	return TEEX_SUCCESS;
}

teex_status_t generateMsgToWorker(teexTask *taskInfo, char **msgToWorker)
{
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (msgToWorker == NULL)
		return TEEX_ERROR_INTERNAL;
	
	teex_status_t retval = TEEX_ERROR_UNIMPLEMENTED;
	cJSON *msgJson = NULL, *dataArray = NULL, *dataItem = NULL;
	char *msgBody = NULL, *msg = NULL;
	char *input_base64 = NULL, *price = NULL;
	char *inputHash = NULL;
	int i = 0;

	if ((input_base64 = (char*)malloc(B64_ENCODE_LEN(strlen(taskInfo->input) + 1))) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(input_base64, 0, B64_ENCODE_LEN(strlen(taskInfo->input)) + 1);
	base64_encode(taskInfo->input, strlen(taskInfo->input), input_base64);

	if ((inputHash = sha256AndToHexStr(taskInfo->input, strlen(taskInfo->input))) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}

	if ((price = (char*)malloc(65)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(price, 0, 65);
	memset(price, '0', 64);
	strncpy(price + 64 - strlen(taskInfo->price), taskInfo->price, strlen(taskInfo->price));
	
	if ((msgJson = cJSON_CreateObject()) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}

	cJSON_AddStringToObject(msgJson, "task_id", taskInfo->taskID);
	cJSON_AddStringToObject(msgJson, "service_id", taskInfo->serviceID);
	cJSON_AddStringToObject(msgJson, "input", input_base64);
	cJSON_AddStringToObject(msgJson, "data_hash", inputHash);
	cJSON_AddStringToObject(msgJson, "pub_key", taskInfo->publicKey);
	cJSON_AddStringToObject(msgJson, "price", price);
	cJSON_AddStringToObject(msgJson, "type", "TEEX");

	if (taskInfo->dataNumber > 0) {
		if ((dataArray = cJSON_CreateArray()) == NULL) {
			retval = TEEX_ERROR_OUT_MEMORY;
			goto out;
		}

		for (i = 0; i < taskInfo->dataNumber; i++) {
			if ((dataItem = cJSON_CreateObject()) == NULL) {
				retval = TEEX_ERROR_OUT_MEMORY;
				goto out;
			}
			cJSON_AddStringToObject(dataItem, "data_id", taskInfo->dataList[i]);
			cJSON_AddItemToArray(dataArray, dataItem);
		}

		cJSON_AddItemToObject(msgJson, "data_list", dataArray);
	}

	msgBody = cJSON_Print(msgJson);

	if ((msg = (char*)malloc(strlen(msgBody) + 128)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(msg, 0, strlen(msgBody) + 128);

	sprintf(msg, "User\n");
	sprintf(msg, "%sContent-Length: %d\n", msg, strlen(msgBody));
	sprintf(msg, "%s%s", msg, msgBody);

	*msgToWorker = msg;
	retval = TEEX_SUCCESS;

out:
	if (input_base64)
		free(input_base64);
	if (price)
		free(price);
	if (msgBody)
		free(msgBody);

	if (msgJson)
		cJSON_Delete(msgJson);
	
	return retval;
}

teex_status_t generateMsgToDispatcher(teexTask *taskInfo, char **msgToDispatcher)
{
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (msgToDispatcher == NULL)
		return TEEX_ERROR_INTERNAL;
	
	cJSON *msgJson = NULL;
	char *msgBody = NULL, *msg = NULL;

	if ((msgJson = cJSON_CreateObject()) == NULL)
		return TEEX_ERROR_OUT_MEMORY;

	cJSON_AddStringToObject(msgJson, "task_id", taskInfo->taskID);
	cJSON_AddStringToObject(msgJson, "service_id", taskInfo->serviceID);
	cJSON_AddStringToObject(msgJson, "pk", taskInfo->publicKey);
	cJSON_AddStringToObject(msgJson, "price", taskInfo->price);

	if ((msgBody = strdup(cJSON_Print(msgJson))) == NULL) {
		free(msgJson);
		return TEEX_ERROR_OUT_MEMORY;
	}
	free(msgJson);

	if ((msg = (char*)malloc(strlen(msgBody) + 128)) == NULL) {
		free(msgBody);
		return TEEX_ERROR_OUT_MEMORY;
	}
	memset(msg, 0, strlen(msgBody) + 128);

	sprintf(msg, "Dispatch\n");
	sprintf(msg, "%sContent-Length: %d\n", msg, strlen(msgBody));
	sprintf(msg, "%s%s", msg, msgBody);
	free(msgBody);

	*msgToDispatcher = msg;

	return TEEX_SUCCESS;
}

teex_status_t getMsgFromDispatcher(TEEX_SSL *ts, char **msgFromDispatcher)
{
	if (ts == NULL || msgFromDispatcher == NULL)
		return TEEX_ERROR_INTERNAL;
	
	char buf[1024], *msg = NULL;
	int msgLen = 0;

	memset(buf, 0, 1024);
	if (teexGetLine(ts, buf, 1024) <= 0) {
		return TEEX_ERROR_NETWORK;
	}

	memset(buf, 0, 1024);
	if (teexGetLine(ts, buf, 1024) <= 0) {
		return TEEX_ERROR_NETWORK;
	}

	if (sscanf(buf, "Content-Length: %d", &msgLen) != 1) {
		return TEEX_ERROR_INVALID_RESP_FORMAT;
	}
	
	if ((msg = (char*)malloc(msgLen + 1)) == NULL)
		return TEEX_ERROR_OUT_MEMORY;
	memset(msg, 0, msgLen + 1);
	
	if (TEEX_readn(ts, msg, msgLen) <= 0) {
		free(msg);
		return TEEX_ERROR_NETWORK;
	}
	

	*msgFromDispatcher = msg;

	return TEEX_SUCCESS;
}

teex_status_t getMsgFromWorker(TEEX_SSL *ts, char **msgFromWorker)
{
	if (ts == NULL || msgFromWorker == NULL)
		return TEEX_ERROR_INTERNAL;
	
	char buf[1024], *msg = NULL;
	int msgLen = 0;

	if (teexGetLine(ts, buf, 1024) <= 0) {
		return TEEX_ERROR_NETWORK;
	}

	if (sscanf(buf, "Content-Length: %d", &msgLen) != 1) {
		return TEEX_ERROR_INVALID_RESULT_FORMAT;
	}

	if ((msg = (char*)malloc(msgLen + 1)) == NULL)
		return TEEX_ERROR_OUT_MEMORY;
	memset(msg, 0, msgLen + 1);

	if (TEEX_readn(ts, msg, msgLen) <= 0) {
		return TEEX_ERROR_NETWORK;
	}
	
	*msgFromWorker = msg;

	return TEEX_SUCCESS;
}

teex_status_t getWorkerFromDispatcher(char *dispatcherAddr, int dispatcherPort,
							teexTask *taskInfo, char **workerAddr, int *workerPort)
{
	/* check the arguments */
	if (dispatcherAddr == NULL)
		return TEEX_ERROR_INVALID_DSPT_ADDR;
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (workerAddr == NULL)
		return TEEX_ERROR_INVALID_WORKER_BUF;
	if (workerPort == NULL)
		return TEEX_ERROR_INVALID_WORKER_BUF;

	
	teex_status_t retval = TEEX_ERROR_UNIMPLEMENTED;
	int dispatcherFD = -1, tmpFD = -1;
	struct timeval timeout = {600, 0};
	TEEX_SSL *ts = NULL;

	char *home = NULL, *cert = NULL, *key = NULL;
	char *msgToDispatcher = NULL, *msgFromDispatcher = NULL;

	cJSON *item = NULL, *workerJson = NULL;
	char *addr = NULL;
	int port = 0;

	/* connect to the dispatcher */
	if ((dispatcherFD = teexConnect(dispatcherAddr, dispatcherPort)) < 0) {
		return TEEX_ERROR_CONNECT_DISPATCHER;
	}

	if (setsockopt(dispatcherFD, SOL_SOCKET, SO_RCVTIMEO, 
					(char*)&timeout, sizeof(struct timeval)) < 0) {
		retval = TEEX_ERROR_NETWORK;
		goto out;
	}

	if ((ts = TEEX_SSL_new(TEEX_SSL_SINGLE_RA_INITIATOR)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	TEEX_SSL_set_fd(ts, dispatcherFD);

	/* get the default client cert */
	if ((home = getpwuid(getuid())->pw_dir) == NULL) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	
	if ((cert = (char*)malloc(strlen(home) + 128)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(cert, 0, strlen(home)+128);
	snprintf(cert, strlen(home)+128, "%s/%s", home, ".teex/.private/client.crt");

	if ((key = (char*)malloc(strlen(home) + 128)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(key, 0, strlen(home)+128);
	snprintf(key, strlen(home)+128, "%s/%s", home, ".teex/.private/client.key");

	if ((tmpFD=open(cert, O_RDONLY)) < 0) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	close(tmpFD);

	if ((tmpFD = open(key, O_RDONLY)) < 0) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	close(tmpFD);
	tmpFD = 0;

	/* setup the TEEX_SSL connection */
	TEEX_SSL_set_ias_client_cert_key(ts, cert, key);
	TEEX_SSL_set_attributes(ts, 0);
	
	if (TEEX_connect(ts) != 0) {
		retval = TEEX_ERROR_ATTESTATION_FAILED;
		goto out;
	}

	/* send the message to the dispatcher */
	if ((retval=generateMsgToDispatcher(taskInfo, &msgToDispatcher)) != TEEX_SUCCESS) {
		goto out;
	}


	if (TEEX_writen(ts, msgToDispatcher, strlen(msgToDispatcher)) <= 0) {
		retval = TEEX_ERROR_NETWORK;
		goto out;
	}

	/* get the return message from the dispatcher */
	if ((retval=getMsgFromDispatcher(ts, &msgFromDispatcher)) != TEEX_SUCCESS) {
		goto out;
	}

	workerJson = cJSON_Parse(msgFromDispatcher);

	if (workerJson == NULL) {
		retval = TEEX_ERROR_INVALID_JSON;
		goto out;
	}

	if ((item = cJSON_GetObjectItem(workerJson, "port")) != NULL
		&& cJSON_IsNumber(item)) {
		port = item->valueint;
	} else {
		retval = TEEX_ERROR_INVALID_JSON;
		goto out;
	}


	if ((item = cJSON_GetObjectItem(workerJson, "address")) != NULL
		&& cJSON_IsString(item)) {
		addr = strdup(item->valuestring);
	} else {
		retval = TEEX_ERROR_INVALID_JSON;
		goto out;
	}

	*workerAddr = addr;
	*workerPort = port;
	retval = TEEX_SUCCESS;

out:
	/* free all pointers */
	if (ts)
		TEEX_SSL_free(ts);
	if (workerJson)
		cJSON_Delete(workerJson);

	/* close all FDs */
	if (dispatcherFD > 0)
		close(dispatcherFD);
	if (tmpFD > 0)
		close(tmpFD);

	/* free all string */
	if (cert)
		free(cert);
	if (key)
		free(key);
	if (msgToDispatcher)
		free(msgToDispatcher);
	if (msgFromDispatcher)
		free(msgFromDispatcher);
	
	return retval;
}




teex_status_t getResultFromWorker(char *workerAddr, int workerPort,
									teexTask *taskInfo, char **returnMsg,
									int *msgLen)
{
	teex_status_t ret = TEEX_ERROR_UNIMPLEMENTED;
	
	/* check the arguments */
	if (workerAddr == NULL)
		return TEEX_ERROR_INVALID_WORKER_ADDR;
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (returnMsg == NULL || msgLen == NULL)
		return TEEX_ERROR_INVALID_RESULT_BUF;

	
	teex_status_t retval = TEEX_ERROR_UNIMPLEMENTED;
	int workerFD = -1, tmpFD = -1;
	struct timeval timeout = {600, 0};
	TEEX_SSL *ts = NULL;

	char *home = NULL, *cert = NULL, *key = NULL;
	char *msgToWorker = NULL, *msgFromWorker = NULL;
	char *result = NULL;
	int resultLen = 0;

	cJSON *item = NULL, *msgJson = NULL;

	/* connect to the dispatcher */
	if ((workerFD = teexConnect(workerAddr, workerPort)) < 0) {
		return TEEX_ERROR_CONNECT_WORKER;
	}

	if (setsockopt(workerFD, SOL_SOCKET, SO_RCVTIMEO, 
					(char*)&timeout, sizeof(struct timeval)) < 0) {
		retval = TEEX_ERROR_NETWORK;
		goto out;
	}

	if ((ts = TEEX_SSL_new(TEEX_SSL_SINGLE_RA_INITIATOR)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	TEEX_SSL_set_fd(ts, workerFD);

	/* get the default client cert */
	if ((home = getpwuid(getuid())->pw_dir) == NULL) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	
	if ((cert = (char*)malloc(strlen(home) + 128)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(cert, 0, strlen(home)+128);
	snprintf(cert, strlen(home)+128, "%s/%s", home, ".teex/.private/client.crt");

	if ((key = (char*)malloc(strlen(home) + 128)) == NULL) {
		retval = TEEX_ERROR_OUT_MEMORY;
		goto out;
	}
	memset(key, 0, strlen(home)+128);
	snprintf(key, strlen(home)+128, "%s/%s", home, ".teex/.private/client.key");

	if ((tmpFD=open(cert, O_RDONLY)) < 0) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	close(tmpFD);

	if ((tmpFD = open(key, O_RDONLY)) < 0) {
		retval = TEEX_ERROR_INVALID_CLIENT_CERT;
		goto out;
	}
	close(tmpFD);
	tmpFD = 0;
	
	/* setup the TEEX_SSL connection */
	TEEX_SSL_set_ias_client_cert_key(ts, cert, key);
	TEEX_SSL_set_attributes(ts, 0);
	
	if (TEEX_connect(ts) != 0) {
		retval = TEEX_ERROR_ATTESTATION_FAILED;
		goto out;
	}

	
	if ((retval = generateMsgToWorker(taskInfo, &msgToWorker)) != TEEX_SUCCESS) {
		goto out;
	}

	if (TEEX_writen(ts, msgToWorker, strlen(msgToWorker)) <= 0) {
		retval = TEEX_ERROR_NETWORK;
		goto out;
	}

	if ((retval = getMsgFromWorker(ts, &msgFromWorker)) != TEEX_SUCCESS) {
		goto out;
	}


	/* parse the msg and get the result */	
	if ((msgJson = cJSON_Parse(msgFromWorker)) == NULL) {
		retval = TEEX_ERROR_INVALID_RESULT_FORMAT;
		goto out;
	}

	if ((item = cJSON_GetObjectItem(msgJson, "err_no")) != NULL
		&& cJSON_IsNumber(item)) {
		if (item->valueint == 0) {
			/* execution succeed, return the result */
			if ((item = cJSON_GetObjectItem(msgJson, "ret_msg")) != NULL
				&& cJSON_IsString(item)) {
				
				resultLen = B64_DECODE_LEN(strlen(item->valuestring)) + 1;
				if ((result = (char*)malloc(resultLen)) == NULL) {
					retval = TEEX_ERROR_OUT_MEMORY;
					goto out;
				}
				memset(result, 0, resultLen);
				resultLen = base64_decode(item->valuestring, result);
				result[resultLen] = 0;
				
				
				*returnMsg = result;
				*msgLen = strlen(result);
				retval = TEEX_SUCCESS;
				goto out;

			} else {
				retval = TEEX_ERROR_INVALID_RESULT_FORMAT;
				goto out;
			}

		} else {
			/* execution failed, return the error message */
			if ((item = cJSON_GetObjectItem(msgJson, "err_msg")) != NULL
				&& cJSON_IsString(item)) {
				
				*returnMsg = strdup(item->valuestring);
				*msgLen = strlen(item->valuestring);

				retval = TEEX_ERROR_WORKER_EXECUTION;
				goto out;
			} else {
				retval = TEEX_ERROR_INVALID_RESULT_FORMAT;
				goto out;
			}
		}

	} else {
		retval = TEEX_ERROR_INVALID_RESULT_FORMAT;
		goto out;
	}

out:
	if (key)
		free(key);
	if (cert)
		free(cert);
	if (msgToWorker)
		free(msgToWorker);
	if (msgFromWorker)
		free(msgFromWorker);
	
	if (ts)
		TEEX_SSL_free(ts);
	if (msgJson)
		cJSON_Delete(msgJson);
	
	if (workerFD)
		close(workerFD);
	if (tmpFD)
		close(tmpFD);
	
	return retval;
}



teex_status_t runTask(char *dispatcherAddr, int dispatcherPort,
							teexTask *taskInfo, char **returnMsg,
							int *msgLen)
{

	teex_status_t ret = TEEX_ERROR_UNIMPLEMENTED;
	int workerPort = 0;
	char *workerAddr;

	if (dispatcherAddr == NULL)
		return TEEX_ERROR_INVALID_DSPT_ADDR;
	if (taskInfo == NULL)
		return TEEX_ERROR_INVALID_TASKINFO;
	if (returnMsg == NULL)
		return TEEX_ERROR_INVALID_RESULT_BUF;
	if (msgLen == NULL)
		return TEEX_ERROR_INVALID_SIZE_BUF;

	if ((ret = getWorkerFromDispatcher(dispatcherAddr, dispatcherPort, 
				taskInfo, &workerAddr, &workerPort)) != 0) {
		return ret;
	}

	if ((ret = getResultFromWorker(workerAddr, workerPort, taskInfo, returnMsg, msgLen)) != 0) {
		return ret;
	}

	return TEEX_SUCCESS;
}
